// Copyright (c) 2025 Shenzhen Kaihong Digital Industry Development Co., Ltd.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::{collections::{hash_map, HashMap}, fmt::Debug, ops::Range, path::Path, ptr::NonNull};

use fixedbitset::FixedBitSet;
use petgraph::{algo::Cycle, dot::{Config, Dot}, graph::{EdgeIndex, EdgeReference, Edges, Neighbors, NodeIndex}, visit::{Reversed, Visitable}, Direction, Graph};
use sha2::{Digest, Sha256};

use crate::{config::{pkg_config::{PkgConfig, PkgConfigLoadError}, system::{toolchain::SubToolchainKind, OverlayConfig}, PkgDepConfig}, database::local::PKGS_CONFIG_DIR, Arch, PkgSpec, Version};

use super::DepKind;

#[derive(Debug, Clone)]
pub struct DepNode {
    pub spec: PkgSpec,
    pub config: PkgConfig,
}

#[derive(Debug)]
pub struct DepEdge {
    pub kind: DepKind,
    pub config: Option<NonNull<PkgDepConfig>>,
}

#[derive(Default, Debug)]
pub struct DepGraph {
    pub(super) graph: Graph<DepNode, DepEdge>,
    node_map: HashMap<String, HashMap<Version, Vec<NodeIndex>>>,
    pub(super) edge_map: HashMap<(NodeIndex, NodeIndex, DepKind), EdgeIndex>,
    pub id_map: DepIdMap,
}

#[derive(Default, Debug)]
pub struct DepIdMap {
    pub subtoolchain_ranges: HashMap<SubToolchainKind, ToolchainRanges>,
    pub temp_toolchain_range: Range<usize>,
    pub(super) toolchain_specs: FixedBitSet,
    pub(super) non_toolchain_specs_range: Range<usize>,
    pub preload_count: usize,
    pub(super) deploy_runtime_depend_specs: FixedBitSet,
}

#[derive(Default, Debug)]
pub struct ToolchainRanges {
    pub host_toolchain_range: Range<usize>,
    pub cross_toolchain_ranges: HashMap<Arch, Range<usize>>,
}

unsafe impl Send for DepEdge {}

unsafe impl Sync for DepEdge {}

impl DepGraph {
    pub fn dot(&self) -> Dot<'_, &Graph<DepNode, DepEdge>> {
        fn get_node_attributes(_graph: &Graph<DepNode, DepEdge>, node_ref: (NodeIndex, &DepNode)) -> String {
            let (id, node) = node_ref;
            format!("label=\"{}({})\"", node.spec, id.index())
        }

        fn get_edge_attributes(_graph: &Graph<DepNode, DepEdge>, edge_ref: EdgeReference<'_, DepEdge>) -> String {
            let edge = edge_ref.weight();
            format!("label=\"{}\"", edge.kind)
        }

        Dot::with_attr_getters(
            &self.graph,
            &[Config::NodeNoLabel, Config::EdgeNoLabel],
            &get_edge_attributes,
            &get_node_attributes,
        )
    }

    pub fn reversed(&self) -> Reversed<&Graph<DepNode, DepEdge>> {
        Reversed(&self.graph)
    }

    pub fn visit_map(&self) -> FixedBitSet {
        self.graph.visit_map()
    }

    pub fn toposort(&mut self) -> Result<Vec<NodeIndex>, Cycle<NodeIndex>> {
        petgraph::algo::toposort(&self.graph, None)
    }

    pub fn node_count(&self) -> usize {
        self.graph.node_count()
    }

    pub fn is_leaf(&self, id: NodeIndex) -> bool {
        self.graph.first_edge(id, Direction::Outgoing).is_none()
    }

    pub fn node(&self, id: NodeIndex) -> Option<&DepNode> {
        self.graph.node_weight(id)
    }

    pub fn dep_nodes(&self, id: NodeIndex) -> Neighbors<'_, DepEdge> {
        self.graph.neighbors_directed(id, Direction::Incoming)
    }

    pub fn dep_edges(&self, id: NodeIndex) -> Edges<'_, DepEdge, petgraph::Directed> {
        self.graph.edges_directed(id, Direction::Incoming)
    }

    pub fn build_dep_edges(&self, id: NodeIndex) -> impl Iterator<Item = EdgeReference<DepEdge>> {
        self.dep_edges(id).filter(|edge| edge.weight().kind == DepKind::Build)
    }

    pub fn host_build_dep_edges(&self, id: NodeIndex) -> impl Iterator<Item = EdgeReference<DepEdge>> {
        self.dep_edges(id).filter(|edge| edge.weight().kind == DepKind::HostBuild)
    }

    pub fn runtime_dep_edges(&self, id: NodeIndex) -> impl Iterator<Item = EdgeReference<DepEdge>> {
        self.dep_edges(id).filter(|edge| edge.weight().kind == DepKind::Runtime)
    }

    pub async fn add_node(&mut self, spec: PkgSpec, overlay: Option<&OverlayConfig>) -> Result<NodeIndex, PkgConfigLoadError> {
        debug_assert!(spec.features.is_sorted());

        match self.node_map.entry(spec.pkg.name.clone()) {
            hash_map::Entry::Occupied(mut occupied_entry) => {
                match occupied_entry.get_mut().entry(spec.pkg.version.clone()) {
                    hash_map::Entry::Occupied(mut occupied_entry) => {
                        let ids = occupied_entry.get_mut();
                        let id = Self::find_match_node(&self.graph, ids, &spec);
                        if id != NodeIndex::end() {
                            return Ok(id);
                        }

                        let config = PkgConfig::load(&spec, overlay).await?;
                        let id = Self::add_node_inner(&mut self.graph, spec, config);
                        ids.push(id);
                        Ok(id)
                    },
                    hash_map::Entry::Vacant(vacant_entry) => {
                        let config = PkgConfig::load(&spec, overlay).await?;
                        let id = Self::add_node_inner(&mut self.graph, spec, config);
                        vacant_entry.insert(vec![id]);
                        Ok(id)
                    },
                }
            },
            hash_map::Entry::Vacant(vacant_entry) => {
                let config  = PkgConfig::load(&spec, overlay).await?;
                let version = spec.pkg.version.clone();
                let id = Self::add_node_inner(&mut self.graph, spec, config);

                let mut submap = HashMap::new();
                submap.insert(version, vec![id]);
                vacant_entry.insert(submap);

                Ok(id)
            },
        }
    }

    pub(super) fn add_edge(&mut self, from: NodeIndex, to: NodeIndex, edge: DepEdge) -> EdgeIndex {
        Self::add_edge_inner(&mut self.edge_map, &mut self.graph, from, to, edge)
    }

    pub(super) fn add_edge_inner(
        edge_map: &mut HashMap<(NodeIndex, NodeIndex, DepKind), EdgeIndex>,
        graph: &mut Graph<DepNode, DepEdge>,
        from: NodeIndex,
        to: NodeIndex,
        edge: DepEdge,
    ) -> EdgeIndex {
        match edge_map.entry((from, to, edge.kind)) {
            hash_map::Entry::Occupied(occupied_entry) => {
                *occupied_entry.get()
            },
            hash_map::Entry::Vacant(vacant_entry) => {
                let id = graph.add_edge(from, to, edge);
                vacant_entry.insert(id);
                id
            },
        }
    }

    fn add_node_inner(graph: &mut Graph<DepNode, DepEdge>, spec: PkgSpec, config: PkgConfig) -> NodeIndex {
        graph.add_node(DepNode {
            spec,
            config,
        })
    }

    fn find_match_node(graph: &Graph<DepNode, DepEdge>, ids: &[NodeIndex], spec: &PkgSpec) -> NodeIndex {
        for &id in ids {
            let node = graph.node_weight(id).unwrap();
            if node.spec.arch == spec.arch && Self::match_features(node, spec) {
                return id;
            }
        }
        NodeIndex::end()
    }

    fn match_features(node: &DepNode, spec: &PkgSpec) -> bool {
        if spec.features.is_empty() {
            node.spec_features() == &node.config.features.default
        } else {
            node.spec_features() == &spec.features
        }
    }
}

impl DepNode {
    pub fn spec_features(&self) -> &Vec<String> {
        if self.spec.features.is_empty() {
            &self.config.features.default
        } else {
            &self.spec.features
        }
    }

    pub async fn spec_hash(&self, hasher: &mut Sha256) -> anyhow::Result<()> {
        hasher.update(&self.spec.pkg.name);
        hasher.update(self.spec.pkg.version.as_str());
        for feature in self.spec_features() {
            hasher.update(feature);
        }
        hasher.update([self.spec.arch as u8]);
        self.config.spec_hash(
            hasher,
            Path::new(PKGS_CONFIG_DIR.as_str()).join(self.spec.pkg.to_string()),
            !self.spec.arch.is_host(),
        ).await?;
        Ok(())
    }
}

impl DepEdge {
    /// # Safety
    ///
    /// The pointer is pointing to a [PkgDepConfig] of the corresonding [DepNode] in the [DepGraph].
    /// The caller must ensure that the pointer is valid and that the [DepNode] is not dropped before the pointer.
    pub unsafe fn get_config(&self) -> Option<&PkgDepConfig> {
        self.config.map(|c| c.as_ref())
    }
}

impl DepIdMap {
    pub fn specs(&self) -> impl Iterator<Item = NodeIndex> + '_ {
        self.toolchain_specs.ones()
            .chain(self.non_toolchain_specs_range.clone())
            .chain(self.deploy_runtime_depend_specs.ones())
            .map(NodeIndex::new)
    }
}